package com.hxb.example.controller;

import com.hxb.example.service.ChatModelService;
import com.hxb.example.util.PlayWrightUtil;
import com.microsoft.playwright.Browser;
import com.microsoft.playwright.BrowserType;
import com.microsoft.playwright.ElementHandle;
import com.microsoft.playwright.Page;
import com.microsoft.playwright.Playwright;
import jakarta.annotation.Resource;
import jakarta.servlet.http.HttpServletResponse;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.function.Supplier;

@Slf4j
@RestController
public class EmbeddingController {

    @Resource
    private ChatModelService chatModelService;
    @Resource
    private HttpServletResponse response;

    // @GetMapping("/ai/embedding")
    // public Map embed(@RequestParam(value = "message", defaultValue = "Tell me a joke") String message) {
    //     EmbeddingResponse embeddingResponse = this.embeddingModel.embedForResponse(List.of(message));
    //     return Map.of("embedding", embeddingResponse);
    // }

    @GetMapping("/ai/call")
    public Object call(@RequestParam(value = "message", defaultValue = "Tell me a joke") String message) {
        Object response = this.chatModelService.call(message);
        return Map.of("embedding", response);
    }

    @GetMapping("/ai/stream/call")
    public SseEmitter streamCall(@RequestParam(value = "message", defaultValue = "Tell me a joke") String message) {
        return response(() -> this.chatModelService.streamCall(message));
    }

    @GetMapping("/ai/function/call")
    public Object functionCall(@RequestParam(value = "message", defaultValue = "Tell me a joke") String message) {
        return this.chatModelService.functionCall(message);
    }

    @GetMapping("/ai/function/stream/call")
    public Object functionStreamCall(@RequestParam(value = "message", defaultValue = "Tell me a joke") String message) {
        return response(() -> this.chatModelService.functionStreamCall(message));
    }

    @GetMapping("playwright")
    public Object playwright() {
        try (Page page = PlayWrightUtil.initDefaultPage()) {
            page.navigate("https://www.toutiao.com/");
            // home-hotboard
            List<ElementHandle> tableRows = page.querySelectorAll(".news-title");
            for (ElementHandle row : tableRows) {
                System.out.println(row.textContent());
            }
        } catch (Exception e) {
            log.error("趴取头条热点新闻异常", e);
        }
        return null;
    }

    private SseEmitter response(Supplier<Flux<ChatResponse>> supplier) {
        response.setContentType("text/event-stream");
        response.setCharacterEncoding("UTF-8");
        SseEmitter emitter = new SseEmitter(300_000L);
        Flux<ChatResponse> stream = supplier.get();
        if (null == stream) {
            return null;
        }
        stream.doOnNext(chatResponse -> {
                    try {
                        emitter.send(chatResponse.getResult().getOutput().getContent());
                    } catch (IOException e) {
                        log.error("SEE推送消息失败, 失败消息: {}", chatResponse.getResult().getOutput().getContent(), e);
                        emitter.completeWithError(e);
                    }
                })
                .doOnError(emitter::completeWithError)
                .doOnComplete(emitter::complete).subscribe();
        return emitter;
    }
}

